{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transformers installation\n",
    "! pip install transformers datasets evaluate accelerate\n",
    "# To install from source instead of the last release, comment the command above and uncomment the following one.\n",
    "# ! pip install git+https://github.com/huggingface/transformers.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Keypoint Detection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Keypoint detection identifies and locates specific points of interest within an image. These keypoints, also known as landmarks, represent meaningful features of objects, such as facial features or object parts. These models take an image input and return the following outputs:\n",
    "\n",
    "- **Keypoints and Scores**: Points of interest and their confidence scores.\n",
    "- **Descriptors**: A representation of the image region surrounding each keypoint, capturing its texture, gradient, orientation and other properties.\n",
    "\n",
    "In this guide, we will show how to extract keypoints from images.\n",
    "\n",
    "For this tutorial, we will use [SuperPoint](https://huggingface.co/docs/transformers/main/en/tasks/./model_doc/superpoint), a foundation model for keypoint detection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoImageProcessor, SuperPointForKeypointDetection\n",
    "processor = AutoImageProcessor.from_pretrained(\"magic-leap-community/superpoint\")\n",
    "model = SuperPointForKeypointDetection.from_pretrained(\"magic-leap-community/superpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test the model on the images below.\n",
    "\n",
    "<div style=\"display: flex; align-items: center;\">\n",
    "    <img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg\"\n",
    "         alt=\"Bee\"\n",
    "         style=\"height: 200px; object-fit: contain; margin-right: 10px;\">\n",
    "    <img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png\"\n",
    "         alt=\"Cats\"\n",
    "         style=\"height: 200px; object-fit: contain;\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "import requests\n",
    "import cv2\n",
    "\n",
    "\n",
    "url_image_1 = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg\"\n",
    "image_1 = Image.open(requests.get(url_image_1, stream=True).raw)\n",
    "url_image_2 = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png\"\n",
    "image_2 = Image.open(requests.get(url_image_2, stream=True).raw)\n",
    "\n",
    "images = [image_1, image_2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now process our inputs and infer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = processor(images,return_tensors=\"pt\").to(model.device, model.dtype)\n",
    "outputs = model(**inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model output has relative keypoints, descriptors, masks and scores for each item in the batch. The mask highlights areas of the image where keypoints are present."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SuperPointKeypointDescriptionOutput(loss=None, keypoints=tensor([[[0.0437, 0.0167],\n",
    "         [0.0688, 0.0167],\n",
    "         [0.0172, 0.0188],\n",
    "         ...,\n",
    "         [0.5984, 0.9812],\n",
    "         [0.6953, 0.9812]]]), \n",
    "         scores=tensor([[0.0056, 0.0053, 0.0079,  ..., 0.0125, 0.0539, 0.0377],\n",
    "        [0.0206, 0.0058, 0.0065,  ..., 0.0000, 0.0000, 0.0000]],\n",
    "       grad_fn=<CopySlices>), descriptors=tensor([[[-0.0807,  0.0114, -0.1210,  ..., -0.1122,  0.0899,  0.0357],\n",
    "         [-0.0807,  0.0114, -0.1210,  ..., -0.1122,  0.0899,  0.0357],\n",
    "         [-0.0807,  0.0114, -0.1210,  ..., -0.1122,  0.0899,  0.0357],\n",
    "         ...],\n",
    "       grad_fn=<CopySlices>), mask=tensor([[1, 1, 1,  ..., 1, 1, 1],\n",
    "        [1, 1, 1,  ..., 0, 0, 0]], dtype=torch.int32), hidden_states=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To plot actual keypoints in the image, we need to postprocess the output. To do so, we have to pass the actual image sizes to `post_process_keypoint_detection` along with outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_sizes = [(image.size[1], image.size[0]) for image in images]\n",
    "outputs = processor.post_process_keypoint_detection(outputs, image_sizes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The outputs are now a list of dictionaries where each dictionary is a processed output of keypoints, scores and descriptors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[{'keypoints': tensor([[ 226,   57],\n",
    "          [ 356,   57],\n",
    "          [  89,   64],\n",
    "          ...,\n",
    "          [3604, 3391]], dtype=torch.int32),\n",
    "  'scores': tensor([0.0056, 0.0053, ...], grad_fn=<IndexBackward0>),\n",
    "  'descriptors': tensor([[-0.0807,  0.0114, -0.1210,  ..., -0.1122,  0.0899,  0.0357],\n",
    "          [-0.0807,  0.0114, -0.1210,  ..., -0.1122,  0.0899,  0.0357]],\n",
    "         grad_fn=<IndexBackward0>)},\n",
    "    {'keypoints': tensor([[ 46,   6],\n",
    "          [ 78,   6],\n",
    "          [422,   6],\n",
    "          [206, 404]], dtype=torch.int32),\n",
    "  'scores': tensor([0.0206, 0.0058, 0.0065, 0.0053, 0.0070, ...,grad_fn=<IndexBackward0>),\n",
    "  'descriptors': tensor([[-0.0525,  0.0726,  0.0270,  ...,  0.0389, -0.0189, -0.0211],\n",
    "          [-0.0525,  0.0726,  0.0270,  ...,  0.0389, -0.0189, -0.0211]}]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use these to plot the keypoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "for i in range(len(images)):\n",
    "  keypoints = outputs[i][\"keypoints\"]\n",
    "  scores = outputs[i][\"scores\"]\n",
    "  descriptors = outputs[i][\"descriptors\"]\n",
    "  keypoints = outputs[i][\"keypoints\"].detach().numpy()\n",
    "  scores = outputs[i][\"scores\"].detach().numpy()\n",
    "  image = images[i]\n",
    "  image_width, image_height = image.size\n",
    "\n",
    "  plt.axis('off')\n",
    "  plt.imshow(image)\n",
    "  plt.scatter(\n",
    "      keypoints[:, 0],\n",
    "      keypoints[:, 1],\n",
    "      s=scores * 100,\n",
    "      c='cyan',\n",
    "      alpha=0.4\n",
    "  )\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below you can see the outputs.\n",
    "\n",
    "<div style=\"display: flex; align-items: center;\">\n",
    "    <img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee_keypoint.png\"\n",
    "         alt=\"Bee\"\n",
    "         style=\"height: 200px; object-fit: contain; margin-right: 10px;\">\n",
    "    <img src=\"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats_keypoint.png\"\n",
    "         alt=\"Cats\"\n",
    "         style=\"height: 200px; object-fit: contain;\">\n",
    "</div>"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
